/**
 * @file alloc_hook.cc
 * @brief Implementation of memory allocation hook functions.
 *
 * This file contains the implementation of memory allocation hook functions, which override the default memory
 * allocation functions such as malloc, calloc, realloc, and free.
 */
#include "alloc_hook.h"

#include <dirent.h>
#include <dlfcn.h>
#include <execinfo.h>
#include <fcntl.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <unistd.h>
#include <zlib.h>

#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <memory>
#include <new>
#include <sstream>
#include <string>
#include <vector>

#include "alloc_log.h"

#if __cplusplus >= 201103L || (defined(_MSC_VER) && _MSC_VER >= 1900)
#define CPP_NOTHROW noexcept
#define CPP_BADALLOC
#else
#define CPP_NOTHROW throw()
#define CPP_BADALLOC throw(std::bad_alloc)
#endif

constexpr int64_t kListTotalLen = 20000;
constexpr int32_t kListNum = 200;
constexpr int32_t kListLen = kListTotalLen / kListNum;

extern "C" void *__libc_malloc(size_t size);
extern "C" void *__libc_calloc(size_t nmemb, size_t size);
extern "C" void *__libc_realloc(void *ptr, size_t size);
extern "C" void __libc_free(void *ptr);
static bool backtrace_init = false;            // flag to indicate if backtrace is initialized
static AllocRecord *list_head[kListNum];       // head of the allocation record list
static pthread_spinlock_t spinlock[kListNum];  // spinlock for the allocation record list
static ThreadRecord *thread_info = nullptr;    // head of the thread information list
static const char *file_dir = "/tmp";          // directory to store the dump files

static void DumpSignalHandler(int signal_number);
static void AllocRecordInit();

/**
 * @brief Retrieves the ThreadRecord entry for a given process ID.
 *
 * This function searches for the ThreadRecord entry with the specified process ID in the global thread_info linked
 * list. If the entry is found, it is returned. If not found, a new ThreadRecord entry is created and added to the
 * linked list.
 *
 * @param pid The process ID for which to retrieve the ThreadRecord entry.
 * @return A pointer to the ThreadRecord entry for the specified process ID.
 */
static ThreadRecord *GetThreadEntry(int pid) {
  ThreadRecord *entry = nullptr;
  ThreadRecord *p = thread_info;
  if (pid) {
    entry = thread_info;
    while (entry) {
      if (entry->pid == pid) {
        return entry;
      }
      entry = entry->next;
    }
  }

  if (entry == nullptr) {
    entry = static_cast<ThreadRecord *>(__libc_calloc(1, sizeof(ThreadRecord)));
    entry->next = thread_info;
    while (__sync_bool_compare_and_swap(&thread_info, p, entry) == 0) {
      p = thread_info;
      entry->next = thread_info;
    }
  } else {
    // do nothing
  }
  return entry;
}

/**
 * @class MemOverrideInit
 * @brief Initializes memory override functionality.
 *
 * This class is responsible for initializing the memory override functionality. It sets up the signal handler for
 * dumping memory, retrieves the dump directory from the environment variables, initializes the stack trace, and
 * initializes the allocation record.
 */
class MemOverrideInit {
 public:
  /**
   * @brief Constructor for MemOverrideInit.
   *
   * sets up the signal handler, retrieves the dump directory, initializes the stack trace, initializes the
   * allocation record
   */
  MemOverrideInit() {
    int signal_number = SIGUSR1;
    const char *signal_number_str = getenv("DUMP_SIGNAL");

    if (signal_number_str != nullptr) {
      signal_number = static_cast<int>(strtol(signal_number_str, nullptr, 10));
    }
    signal(signal_number, DumpSignalHandler);

    if (getenv("DUMP_DIR")) {
      file_dir = getenv("DUMP_DIR");
    }

    void *stack[MAX_STACK_LEN];
    // backtrace will allocate some memory, call it before use ALLOC_RECORD
    backtrace(stack, MAX_STACK_LEN);
    AllocRecordInit();
    LOG_INFO("memory hook init.\n");

    backtrace_init = true;
  }
  ~MemOverrideInit() = default;
};

class ThreadNameRecord {
 public:
  ThreadNameRecord() {
    pid = syscall(SYS_gettid);  // NOLINT
    entry = GetThreadEntry(0);
    entry->pid = pid;
  }
  ~ThreadNameRecord() = default;
  pid_t pid;
  ThreadRecord *entry;
};

#if 1
#define ALLOC_RECORD(_p, _len, _type)                                                                 \
  do {                                                                                                \
    if ((!backtrace_init) || (!(_p))) {                                                                 \
      break;                                                                                          \
    }                                                                                                 \
    AllocRecord *_entry = AllocRecordListGetEntry(_p, _type);                                         \
    if (_entry) {                                                                                     \
      _entry->pid = thread_local_info.pid;                                                            \
      _entry->stack_len = backtrace(_entry->stack, sizeof(_entry->stack) / sizeof(_entry->stack[0])); \
      _entry->p = _p;                                                                                 \
      _entry->len = (_len);                                                                           \
    }                                                                                                 \
  } while (0)

#define ALLOC_RECORD_DEL(_p)        \
  do {                              \
    if (backtrace_init) {           \
      AllocRecordListFreeEntry(_p); \
    }                               \
  } while (0)
#else
#define ALLOC_RECORD(_p, _len)
#define ALLOC_RECORD_DEL(_p)
#endif

static MemOverrideInit mem_override_entry;
static thread_local ThreadNameRecord thread_local_info;

static int AllocRecordListIncrease(int count, int list_num) {
  auto entry = static_cast<AllocRecord *>(__libc_calloc(count, sizeof(AllocRecord)));
  if (entry == nullptr) {
    LOG_ERROR("malloc failed.\n");
    return 0;
  }
  for (int i = 0; i < count - 1; ++i) {
    entry[i].next = &entry[i + 1];
  }
  pthread_spin_lock(&spinlock[list_num]);
  entry[count - 1].next = list_head[list_num];
  list_head[list_num] = entry;
  pthread_spin_unlock(&spinlock[list_num]);
  return count;
}

static inline AllocRecord *AllocRecordListGetEntry(void *address, unsigned char alloc_type) {
  int list_num = static_cast<int>(reinterpret_cast<uint64_t>(address) / 128 % kListNum);
  AllocRecord *p = list_head[list_num];

  while (p) {
    if (!p->in_use) {
      if (__sync_bool_compare_and_swap(&p->in_use, 0, alloc_type)) {
        break;
      }
    }
    p = p->next;
  }
  if (p == nullptr) {
    if (AllocRecordListIncrease(kListLen, list_num) > 0) {
      p = list_head[list_num];
      while (p) {
        if (!p->in_use) {
          if (__sync_bool_compare_and_swap(&p->in_use, 0, alloc_type)) {
            break;
          }
        }
        p = p->next;
      }
    }
  }

  return p;
}

/**
 * Frees an entry in the AllocRecord list based on the given address.
 *
 * @param address The address of the entry to be freed.
 */
static inline void AllocRecordListFreeEntry(void *address) {
  int list_num = static_cast<int>(reinterpret_cast<uint64_t>(address) / 128 % kListNum);
  AllocRecord *p = list_head[list_num];
  while (p) {
    if (p->p == address && p->in_use) {
      p->in_use = 0;
      break;
    }
    p = p->next;
  }
}

static void AllocRecordInit() {
  for (int i = 0; i < kListNum; ++i) {
    pthread_spin_init(&spinlock[i], PTHREAD_PROCESS_SHARED);
    AllocRecordListIncrease(kListLen, i);
  }
}

#if 1
void *operator new(size_t size) CPP_BADALLOC {
  void *p = __libc_malloc(size);
  ALLOC_RECORD(p, size, ALLOC_IN_USE);
  return p;
}

void *operator new[](size_t size) CPP_BADALLOC {
  void *p = __libc_malloc(size);
  ALLOC_RECORD(p, size, ALLOC_IN_USE);
  return p;
}

void *operator new(size_t size, const std::nothrow_t &) CPP_NOTHROW {
  void *p = __libc_malloc(size);
  ALLOC_RECORD(p, size, ALLOC_IN_USE);
  return p;
}

void *operator new[](size_t size, const std::nothrow_t &) CPP_NOTHROW {
  void *p = __libc_malloc(size);
  ALLOC_RECORD(p, size, ALLOC_IN_USE);
  return p;
}

void operator delete(void *p) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

void operator delete[](void *p) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

void operator delete(void *p, const std::nothrow_t &) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

void operator delete[](void *p, const std::nothrow_t &) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

void operator delete(void *p, size_t) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}
void operator delete[](void *p, size_t) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

void operator delete(void *p, size_t, const std::nothrow_t &) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}
void operator delete[](void *p, std::size_t, const std::nothrow_t &) CPP_NOTHROW {
  ALLOC_RECORD_DEL(p);
  __libc_free(p);
}

#endif

extern "C" void *malloc(size_t size) {
  void *p = __libc_malloc(size);
  if (backtrace_init) {
    ALLOC_RECORD(p, size, ALLOC_IN_USE);
  }
  return p;
}

extern "C" void *calloc(size_t nmemb, size_t size) {
  void *p = __libc_calloc(nmemb, size);
  ALLOC_RECORD(p, nmemb * size, ALLOC_IN_USE);
  return p;
}

extern "C" void *realloc(void *ptr, size_t size) {
  ALLOC_RECORD_DEL(ptr);
  void *p = __libc_realloc(ptr, size);
  ALLOC_RECORD(p, size, ALLOC_IN_USE);
  return p;
}

extern "C" void free(void *ptr) {
  ALLOC_RECORD_DEL(ptr);
  __libc_free(ptr);
}

/**
 * @brief Overrides the mmap64 function to record memory allocations.
 *
 * This function is used to allocate memory using the mmap64 system call. It overrides the default
 * mmap64 function and adds a record of the allocated memory using the ALLOC_RECORD macro.
 */
extern "C" void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
  void *p = reinterpret_cast<void *>(syscall(SYS_mmap, addr, length, prot, flags, fd, offset));
  ALLOC_RECORD(p, length, MMAP_IN_USE);
  return p;
}

extern "C" void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
  void *p = reinterpret_cast<void *>(syscall(SYS_mmap, addr, length, prot, flags, fd, offset));
  ALLOC_RECORD(p, length, MMAP_IN_USE);
  return p;
}

extern "C" int munmap(void *start, size_t length) {
  ALLOC_RECORD_DEL(start);
  return static_cast<int>(syscall(SYS_munmap, start, length));
}

/**
 * Compresses a file using gzip compression.
 *
 * @param infile The path to the input file.
 * @param outfile The path to the output file.
 */
void FileCompress(const char *infile, const char *outfile) {
  std::ifstream in(infile, std::ios::binary);
  if (!in) {
    LOG_ERROR("fopen-%s: %m\n", infile);
    return;
  } else {
    // do nothing
  }

  gzFile out = gzopen(outfile, "wb9");
  if (!out) {
    LOG_ERROR("can't gzopen %s\n", outfile);
    return;
  } else {
    // do nothing
  }

  char buf[1024];
  while (in.read(buf, sizeof(buf))) {
    std::streamsize len = in.gcount();
    if (gzwrite(out, buf, static_cast<unsigned>(len)) != len) {
      int err;
      LOG_ERROR("%s\n", gzerror(out, &err));
    } else {
      // do nothing
    }
  }
  if (in.bad()) {
    LOG_ERROR("fread: %m\n");
  } else {
    // do nothing
  }

  in.close();
  gzclose(out);
  unlink(infile);
}

/**
 * @brief Retrieves information about each thread in the current process.
 *
 * This function opens the task directory in the /proc/self directory and iterates through each entry.
 * For each entry, it retrieves the thread PID and opens the stat file for that thread.
 * It then reads the contents of the stat file and parses it to extract thread information such as PID and name.
 * The extracted thread information is stored in a ThreadRecord object and added to the provided vector.
 *
 * @param thread_info_list A reference to a vector of ThreadRecord objects to store the retrieved thread information.
 */
static void GetThreadInfo(std::vector<ThreadRecord> &thread_info_list) {
  // Open the task directory
  DIR *task_dir = opendir("/proc/self/task");
  if (task_dir == nullptr) {
    if (write(STDERR_FILENO, "open task dir failed.\n", sizeof("open task dir failed.\n") - 1) < 0) {
      LOG_ERROR("open task dir failed");
    } else {
      // do nothing
    }
    return;
  } else {
    // do nothing
  }

  const dirent *task_entry = nullptr;
  uint32_t thread_pid = 0;
  char file_name[128] = {};
  char buff[1024] = {};
  int fd = -1;
  char *ptr = nullptr;
  char *comm = nullptr;
  pid_t pid = 0;
  ThreadRecord *thread_entry = nullptr;

  // Iterate through each entry in the task directory
  while ((task_entry = readdir(task_dir)) != nullptr) {
    // Skip non-directory entries and "." or ".." entries
    if (task_entry->d_type != DT_DIR || strcmp(task_entry->d_name, ".") == 0 || strcmp(task_entry->d_name, "..") == 0) {
      continue;
    }

    // Get thread PID and open the stat file for the thread
    thread_pid = static_cast<uint32_t>(atoi(task_entry->d_name));
    snprintf(file_name, sizeof(file_name), "/proc/self/task/%u/stat", thread_pid);
    fd = open(file_name, O_RDONLY);
    if (fd < 0) {
      if (write(STDERR_FILENO, "open stat file failed.\n", sizeof("open stat file failed.\n") - 1) < 0) {
        LOG_ERROR("open stat file failed");
      }
      continue;
    }

    // Read the stat file contents
    if (read(fd, buff, sizeof(buff)) < 0) {
      close(fd);
      continue;
    }
    close(fd);

    // Parse the stat file contents to extract thread information
    ptr = strrchr(buff, ')');
    *ptr = '\0';
    comm = strchr(buff, '(');
    *comm = '\0';

    sscanf(buff, "%lu", reinterpret_cast<uint64_t *>(&pid));
    thread_entry = GetThreadEntry(pid);
    thread_entry->pid = pid;
    strncpy(thread_entry->name, comm + 1, sizeof(thread_entry->name) - 1);
    sscanf(ptr + 2,
           "%*s %*s %*s %*s %*s "
           "%*s %*s %*s %*s %*s "
           "%*s %*s %*s %*s %*s "
           "%*s %*s %*s %*s %*s "
           "%*s %lu %*s %*s %*s "
           "%*s %*s %*s %*s %*s "
           "%*s %*s",
           reinterpret_cast<uint64_t *>(&thread_entry->kstk_esp));

    thread_info_list.push_back(*thread_entry);
  }

  closedir(task_dir);
}

/**
 * @brief Dumps heap information to a file.
 *
 * This function opens the smaps file and the output file specified by `file_name`.
 * It writes the provided `header` to the output file, followed by the contents of the smaps file.
 * Then, it writes the thread information from `thread_info_list` to the output file.
 * Finally, it writes the allocation records to the output file.
 *
 * @param file_name The name of the output file.
 * @param header The header information to be written to the output file.
 * @param thread_info_list The list of thread information to be written to the output file.
 * @return true if the heap information was successfully dumped to the file, false otherwise.
 */
static bool DumpHeapInfo(const char *const file_name, AllocDataHeader &header,
                         const std::vector<ThreadRecord> &thread_info_list) {
  // Open smaps file
  const int smaps_fd = open("/proc/self/smaps", O_RDONLY);
  if (smaps_fd < 0) {
    if (write(STDERR_FILENO, "open smaps failed.\n", sizeof("open smaps failed.\n") - 1) < 0) {
      LOG_ERROR("open smaps failed");
    }
    return false;
  }

  // Open or create the output file
  const int fd = open(file_name, O_RDWR | O_CREAT | O_TRUNC, 0666);
  if (fd < 0) {
    if (write(STDERR_FILENO, "dump heap info failed.\n", sizeof("dump heap info failed.\n") - 1) < 0) {
      LOG_ERROR("dump heap info failed");
    }
    close(smaps_fd);
    return false;
  }

  // Write header to the output file
  if (write(fd, &header, sizeof(header)) < 0) {
    LOG_ERROR("write header failed");
    close(smaps_fd);
    close(fd);
    return false;
  }

  // Read from smaps file and write to output file
  char buff[1024];
  ssize_t len;
  while ((len = read(smaps_fd, buff, sizeof(buff))) > 0) {
    if (write(fd, buff, static_cast<size_t>(len)) < 0) {
      LOG_ERROR("write smaps data failed");
      close(smaps_fd);
      close(fd);
      return false;
    }
    header.smaps_file_len += len;
  }
  close(smaps_fd);

  // Write thread information to output file
  for (const auto &thread_entry : thread_info_list) {
    if (write(fd, &thread_entry, sizeof(ThreadRecord)) < 0) {
      LOG_ERROR("write thread record failed");
      close(fd);
      return false;
    }
    header.thread_info_len += sizeof(ThreadRecord);
  }

  // Write allocation records to output file
  for (int32_t i = 0; i < kListNum; ++i) {
    AllocRecord *p = list_head[i];
    while (p) {
      if (p->in_use) {
        if (write(fd, p, sizeof(AllocRecord)) < 0) {
          LOG_ERROR("write alloc record failed");
          close(fd);
          return false;
        }
        if (p->in_use == ALLOC_IN_USE) {
          header.heap_size += p->len;
        } else if (p->in_use == MMAP_IN_USE) {
          header.mmap_size += p->len;
        }
      }
      header.alloc_list_size += sizeof(AllocRecord);
      p = p->next;
    }
  }

  // Update header with final lengths
  lseek(fd, 0, SEEK_SET);
  if (write(fd, &header, sizeof(header)) < 0) {
    LOG_ERROR("write updated header failed");
    close(fd);
    return false;
  }
  close(fd);
  return true;
}

/**
 * @brief Dumps the signal handler information.
 *
 * This function is responsible for dumping the signal handler information, including heap data, thread information,
 * and allocation records. It opens the task directory in the /proc/self directory to iterate over the threads,
 * retrieves information about each thread, and stores it in the thread_entry structure. It also opens the smaps file
 * in the /proc/self directory to read the memory mapping information. The function then opens a file to write the
 * dumped data, writes the header information, smaps data, thread information, and allocation records to the file,
 * and compresses the file.
 *
 * @param signal_number The signal number.
 */
static void DumpSignalHandler(const int) {
  static int32_t num = 0;
  ++num;

  AllocDataHeader header = {};
  header.version = ALLOC_HOOK_VERSION;
  header.heap_size = 0;
  header.mmap_size = 0;
  header.alloc_list_size = 0;
  header.smaps_file_len = 0;
  header.thread_info_len = 0;

  std::vector<ThreadRecord> thread_info_list;
  GetThreadInfo(thread_info_list);

  char file_name[128] = {};
  char compress_file_name[128] = {};
  snprintf(file_name, sizeof(file_name), "/tmp/heap_raw_data_%d", num);
  snprintf(compress_file_name, sizeof(compress_file_name), "%s/heap_data_%d", file_dir, num);

  if (DumpHeapInfo(file_name, header, thread_info_list)) {
    FileCompress(file_name, compress_file_name);
  } else {
    LOG_ERROR("dump heap info failed.\n");
  }
}
